package rabbit.flt.common.context;

import rabbit.flt.common.Metrics;
import rabbit.flt.common.trace.MethodStackInfo;
import rabbit.flt.common.trace.TraceData;

import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

/**
 * trace context
 */
public class TraceContext {

    private static TraceContext inst = new TraceContext();

    private TraceContext() {
    }

    /**
     * meta上下文
     */
    private ThreadLocal<TraceMetaData> traceMetaContext = new ThreadLocal<>();

    /**
     * 根spanId，对于特定的类方法，该id为父级span Id
     * ！！！ 多线程并发情况下，root spanId可能被污染，不能放在traceMetaContext对象中
     */
    private ThreadLocal<String> rootSpanIdContext = new ThreadLocal<>();

    /**
     * 压栈
     *
     * @param method
     * @return
     */
    public static MethodStackInfo pushStack(Method method) {
        Map<Method, MethodStackInfo> methodStacks = getTraceMetaData().getMethodStacks();
        MethodStackInfo stackInfo = methodStacks.computeIfAbsent(method, m -> {
            MethodStackInfo info = new MethodStackInfo();
            TraceData traceData = new TraceData();
            traceData.setIndex(incrementNode());
            traceData.setTargetMethod(m);
            traceData.setTraceId(getTraceId());
            traceData.setSpanId(calcCurrentSpanId(method));
            traceData.setNodeName(method.getName());
            traceData.setNodeDesc(method.getDeclaringClass().getName() + "." + method.getName());
            traceData.setRequestTime(info.getRequestTime());
            traceData.setServerIp(Metrics.getHostIp());
            traceData.setStatus(TraceData.Status.OK);
            info.setTraceData(traceData);
            return info;
        });
        stackInfo.pushStack();
        return stackInfo;
    }

    public static MethodStackInfo getStackInfo(Method method) {
        return getTraceMetaData().getMethodStacks().get(method);
    }

    /**
     * 获取meta data
     *
     * @return
     */
    public static TraceMetaData getTraceMetaData() {
        return inst.traceMetaContext.get();
    }

    /**
     * 出栈
     *
     * @param method
     * @return
     */
    public static MethodStackInfo popStack(Object method) {
        Map<Method, MethodStackInfo> map = getTraceMetaData().getMethodStacks();
        MethodStackInfo stackInfo = map.get(method);
        if (null == stackInfo) {
            return null;
        }
        stackInfo.popStack();
        if (stackInfo.isPopped()) {
            TraceData traceData = stackInfo.getTraceData();
            if (isTraceOpenedBy(method)) {
                clearContext();
            } else {
                map.remove(method);
                int endIndex = traceData.getSpanId().lastIndexOf('-');
                if (-1 != endIndex) {
                    initRootSpanId(traceData.getSpanId().substring(0, endIndex));
                }
            }
            traceData.updateCost();
        }
        return stackInfo;
    }

    /**
     * 判断是否已经打开trace
     *
     * @return
     */
    public static boolean isTraceOpened() {
        return null != getTraceMetaData();
    }

    public static boolean openTrace(Object target) {
        if (isTraceOpened()) {
            return false;
        }
        openTraceByMeta(new TraceMetaData(target));
        return true;
    }

    public static void openTraceByMeta(TraceMetaData data) {
        inst.traceMetaContext.set(data);
    }

    public static long incrementNode() {
        return getTraceMetaData().incrementNode();
    }

    /**
     * 计算并返回当前节点的span id
     *
     * @param target
     * @return
     */
    public static String calcCurrentSpanId(Object target) {
        String spanId = getRootSpanId();
        if (null == spanId) {
            spanId = "0";
        } else {
            if (isTraceOpenedBy(target)) {
                return spanId;
            }
            spanId = spanId.concat("-").concat(getChildSpanIdIndex(spanId).toString());
        }
        initRootSpanId(spanId);
        return spanId;
    }

    private static Long getChildSpanIdIndex(String spanId) {
        return getSpanIdChildCounter(spanId).getAndAdd(1L);
    }

    /**
     * 获取span id 的counter
     *
     * @param spanId
     * @return
     */
    public static AtomicLong getSpanIdChildCounter(String spanId) {
        return getTraceMetaData().getSpanIdChildCounter(spanId);
    }

    /**
     * 判断是不是当前对象开启的
     *
     * @param target
     * @return
     */
    private static boolean isTraceOpenedBy(Object target) {
        return isTraceOpened() && target.equals(getTraceMetaData().getTraceOpenObject());
    }

    /**
     * 清理context
     */
    public static void clearContext() {
        inst.traceMetaContext.remove();
        inst.rootSpanIdContext.remove();
    }

    /**
     * 设值
     *
     * @param traceData
     */
    public static void setWebTraceData(TraceData traceData) {
        getTraceMetaData().setWebTraceData(traceData);
    }

    public static TraceData getWebTraceData() {
        return getTraceMetaData().getWebTraceData();
    }

    public static void lockWebTraceData() {
        getTraceMetaData().setWebTraceDataLocked(true);
    }

    public static boolean isWebTraceDataLocked() {
        return getTraceMetaData().isWebTraceDataLocked();
    }

    /**
     * 初始化span id，并返回结果
     *
     * @param spanId
     * @return
     */
    public static void initRootSpanId(String spanId) {
        inst.rootSpanIdContext.set(spanId);
    }

    public static String getRootSpanId() {
        return inst.rootSpanIdContext.get();
    }

    /**
     * 获取已经存在的trace id
     *
     * @return
     */
    public static String getTraceId() {
        return getTraceMetaData().getTraceId();
    }

    public static void setTraceId(String traceId) {
        getTraceMetaData().setTraceId(traceId);
    }
}
